﻿using Sagacious.Automation.Host;
using System;
using System.Collections;
using System.Collections.Generic;
using System.Collections.ObjectModel;
using System.Collections.Specialized;
using System.Diagnostics.Contracts;
using System.Globalization;
using System.IO;
using System.Linq;
using System.Management.Automation;
using System.Management.Automation.Host;
using System.Management.Automation.Runspaces;
using System.Reflection;
using System.Text;
using System.Threading.Tasks;
using System.Windows;
using System.Windows.Controls;
using System.Windows.Data;
using System.Windows.Documents;
using System.Windows.Input;
using System.Windows.Media;
using System.Windows.Media.Imaging;
using System.Windows.Navigation;
using System.Windows.Shapes;
using System.Windows.Threading;

namespace Sagacious.Automation.Controls
{
    public interface IPowerShellControl
    {
        ConsoleColor BackgroundColor { get; set; }

        int BufferHeight { get; }

        int BufferWidth { get; }

        /// <summary>
        /// The exit code that the host application will use to exit.
        /// </summary>
        int ExitCode { get; set; }

        ConsoleColor ForegroundColor { get; set; }

        void SetBufferSize(int Width, int Height);

        /// <summary>
        /// Indicator to tell the host application that it should exit.
        /// </summary>
        bool ShouldExit { get; set; }

        void WriteLine(ConsoleColor foregroundColor, ConsoleColor backgroundColor, string value);
    }

    public class PowerShellConsole : RichTextBox, IPowerShellControl
    {

        #region Private Fields

        private BufferCell[,] buffer;

        private string _prompt;

        private Run _commandLine;
        private Run _promptRun;
        private Paragraph _paragraph;

        /// <summary>
        /// Holds a reference to the currently executing pipeline so it can be
        /// stopped by the control-C handler.
        /// </summary>
        private PowerShell currentPowerShell;

        /// <summary>
        /// Used to serialize access to instance data.
        /// </summary>
        private object instanceLock = new object();

        /// <summary>
        /// 
        /// </summary>
        private PowerShellHost powerShellHost;

        /// <summary>
        /// 
        /// </summary>
        private Runspace powerShellRunspace;

        #endregion Private Fields

        #region Constructors

        public PowerShellConsole()
            : base()
        {
            // Set properties of the RichTextbox control to make it look like a console window.
            this.BackgroundColor = ConsoleColor.Blue;
            this.ForegroundColor = ConsoleColor.White;
            this.FontFamily = new System.Windows.Media.FontFamily("Lucida Console");

            // Load the PowerShell host.
            powerShellHost = new PowerShellHost(this);
            powerShellRunspace = RunspaceFactory.CreateRunspace(powerShellHost);
            powerShellRunspace.Open();

            // Create a PowerShell object that will be used to execute the commands
            // to create $profile and load the profiles.
            lock (this.instanceLock)
            {
                this.currentPowerShell = PowerShell.Create();
            }

            try
            {
                this.currentPowerShell.Runspace = this.powerShellRunspace;

                PSCommand[] profileCommands = PSHostUtilities.GetProfileCommands("Sagacious");
                foreach (PSCommand command in profileCommands)
                {
                    this.currentPowerShell.Commands = command;
                    this.currentPowerShell.Invoke();
                }
            }
            finally
            {
                // Dispose of the pipeline line and set it to null, locked because currentPowerShell
                // may be accessed by the ctrl-C handler...
                lock (this.instanceLock)
                {
                    this.currentPowerShell.Dispose();
                    this.currentPowerShell = null;
                }
            }

            // Default prompt. We'll change this later.
            _prompt = @"PS LogViewer:\> ";

            // Create two System.Windows.Documents.Run objects to contain the prompt and the user's
            // command line.
            _promptRun = new Run(_prompt);
            _commandLine = new Run();
            _paragraph = new Paragraph();

            // Create a new command prompt in a new document.
            NewCommandPrompt(true);
        }

        #endregion Constructors

        #region Public Properties

        new private Brush Background
        {
            get
            {
                return base.Background;
            }
            set
            {
                base.Background = value;
            }
        }

        private ConsoleColor backgroundColor;
        public ConsoleColor BackgroundColor
        { 
            get
            {
                return backgroundColor;
            }
            set
            {
                backgroundColor = value;
                this.Background = GetBrushFromConsoleColor(backgroundColor);
            }
        }

        private int bufferHeight = 300;
        public int BufferHeight
        {
            get
            {
                return bufferHeight;
            }
        }

        private int bufferWidth = 80;
        public int BufferWidth
        {
            get
            {
                return bufferWidth;
            }
        }

        new private Brush Foreground
        {
            get
            {
                return base.Foreground;
            }
            set
            {
                base.Foreground = value;
            }
        }

        private ConsoleColor foregroundColor;
        public ConsoleColor ForegroundColor
        {
            get
            {
                return foregroundColor;
            }
            set
            {
                foregroundColor = value;
                this.Foreground = GetBrushFromConsoleColor(foregroundColor);
            }
        }

        /// <summary>
        /// Boolean value telling whether the caret is at the end of the command content.
        /// </summary>
        public bool CommandCaretAtContentEnd
        {
            get
            {
                return CaretPosition.CompareTo(_commandLine.ContentEnd) == 0;
            }
        }

        /// <summary>
        /// Boolean value telling whehter the caret is at the start of the command content.
        /// </summary>
        public bool CommandCaretAtContentStart
        {
            get
            {
                return CaretPosition.CompareTo(_commandLine.ContentStart) == 0;
            }
        }

        /// <summary>
        /// The position of the caret within the command content.
        /// </summary>
        public int CommandCaretPosition
        {
            get
            {
                // TODO: Make it return -1 if the caret is not within the command content.
                return _commandLine.ContentStart.GetOffsetToPosition(CaretPosition);
            }
        }

        /// <summary>
        /// The current command text that will be processed as a PowerShell command.
        /// </summary>
        public string CommandText
        {
            get
            {
                var inlineList = _paragraph.Inlines.ToList();
                var promptIndex = inlineList.IndexOf(_promptRun);
                return inlineList.Where((x, i) => i > promptIndex)
                                 .Cast<Run>()
                                 .Select(x => x.Text)
                                 .Aggregate(string.Empty, (current, part) => current + part);
            }
        }

        /// <summary>
        /// Exit code passed by PowerShell to the console control when it exited.
        /// </summary>
        private int exitCode;
        public int ExitCode
        {
            get
            {
                return this.exitCode;
            }
            set
            {
                this.exitCode = value;
            }
        }

        /// <summary>
        /// Indicator to tell the host application that it should exit.
        /// </summary>
        private bool shouldExit;
        public bool ShouldExit
        {
            get
            {
                return this.shouldExit;
            }
            set
            {
                this.shouldExit = value;
            }
        }

        #endregion Public Properties

        #region Protected Overrides

        protected override void OnKeyUp(KeyEventArgs e)
        {
            base.OnKeyUp(e);

            // If after a key press the _commandLine Run object has been removed from 
            // the document, set the text to an empty string and add it to the end of
            // the document.
            if (_commandLine.PreviousInline == null)
            {
                _commandLine.Text = String.Empty;
                _paragraph.Inlines.Add(_commandLine);
                CaretPosition = _commandLine.ContentEnd;
            }
        }

        /// <summary>
        /// Processes every key pressed when the control has focus.
        /// </summary>
        /// <param name="args">The key pressed arguments.</param>
        protected override void OnPreviewKeyDown(KeyEventArgs args)
        {
            base.OnPreviewKeyDown(args);

            if (args.Key != Key.Tab)
            {
                // Clear auto-completion list
            }

            switch (args.Key)
            {
                case Key.Back:
                    args.Handled = HandleKeyBackspace();
                    break;
                case Key.Delete:
                    args.Handled = HandleKeyDelete();
                    break;
                case Key.Enter:
                    args.Handled = HandleKeyEnter();
                    break;
                case Key.Left:
                    args.Handled = HandleKeyLeft();
                    break;
                case Key.Right:
                    args.Handled = HandleKeyRight();
                    break;
                // Disable up and down navigation keys for now.
                case Key.Down:
                case Key.PageDown:
                case Key.PageUp:
                case Key.Up:
                    args.Handled = true;
                    break;
                default:
                    args.Handled = false;
                    break;
            }
        }

        #endregion Protected Overrides

        /// <summary>
        /// A helper method that builds and executes a pipeline that writes to the default 
        /// output path. Any exceptions that are thrown are just passed to the caller. Since 
        /// all output goes to the default outter, this method() won't return anything.
        /// </summary>
        /// <param name="cmd">
        /// The script to run
        /// </param>
        /// <param name="input">
        /// Any input arguments to pass to the script. If null then nothing is passed in.
        /// </param>
        private void ExecuteCommand(string commandText, object input)
        {
            // Ignore empty command lines.
            if (String.IsNullOrEmpty(commandText))
            {
                return;
            }

            // Create the pipeline object and make it available to the ctrl-C handle through 
            // the currentPowerShell instance variable
            lock (this.instanceLock)
            {
                this.currentPowerShell = PowerShell.Create();
            }

            this.currentPowerShell.Runspace = this.powerShellRunspace;

            try
            {
                this.currentPowerShell.AddScript(commandText);

                // Now add the default outputter to the end of the pipe and indicate that 
                // it should handle both output and errors from the previous commands. This 
                // will result in the output being written using the PSHost and 
                // PSHostUserInterface classes instead of returning objects to the hosting
                // application.
                this.currentPowerShell.AddCommand("out-default");
                this.currentPowerShell.Commands.Commands[0].MergeMyResults(PipelineResultTypes.Error, PipelineResultTypes.Output);

                // If there was any input specified, pass it in, otherwise just
                // execute the pipeline.
                if (input != null)
                {
                    this.currentPowerShell.Invoke(new object[] { input });
                }
                else
                {
                    this.currentPowerShell.Invoke();
                }
            }
            finally
            {
                // Dispose of the pipeline line and set it to null, locked because 
                // currentPowerShell may be accessed by the ctrl-C handler.
                lock (this.instanceLock)
                {
                    this.currentPowerShell.Dispose();
                    this.currentPowerShell = null;
                }
            }
        }

        private Brush GetBrushFromConsoleColor(ConsoleColor color)
        {
            var n = Enum.GetName(typeof(ConsoleColor), color);
            var c = System.Drawing.Color.FromName(n == "DarkYellow" ? "Orange" : n); // bug fix
            return new SolidColorBrush(System.Windows.Media.Color.FromRgb(c.R, c.G, c.B));
        }

        #region Key Handler Functions

        /// <summary>
        /// Handle special situations when the backspace key is pressed to ensure proper
        /// behavior.
        /// </summary>
        /// <returns>
        /// Returns true if this method handled the backspace key press.
        /// </returns>
        private bool HandleKeyBackspace()
        {
            // If the caret is at the begin of the command line, we cant' backspace.
            if (CommandCaretPosition == 0)
            {
                return true;
            }
            // If the caret is in the first position, the RichTextBox will combine the
            // _promptRun and the _commandLine into a single run, so we have to handle
            // the backspace.
            if (CommandCaretPosition == 1)
            {
                _commandLine.Text = _commandLine.Text.Substring(1);
                CaretPosition = _commandLine.ContentStart;
                return true;
            }
            return false;
        }

        /// <summary>
        /// Method used to handle control-C's from the user. It calls the pipeline Stop() 
        /// method to stop execution. If any exceptions occur they are printed to the 
        /// console but otherwise ignored.
        /// </summary>
        /// <param name="sender">
        /// See sender property of ConsoleCancelEventHandler documentation
        /// </param>
        /// <param name="e">
        /// See e property of ConsoleCancelEventHandler documentation
        /// </param>
        private void HandleControlC(object sender, ConsoleCancelEventArgs e)
        {
            try
            {
                lock (this.instanceLock)
                {
                    if (this.currentPowerShell != null && this.currentPowerShell.InvocationStateInfo.State == PSInvocationState.Running)
                    {
                        this.currentPowerShell.Stop();
                    }
                }

                e.Cancel = true;
            }
            catch (Exception exception)
            {
                this.powerShellHost.UI.WriteErrorLine(exception.ToString());
            }
        }

        /// <summary>
        /// Handle special situations when the delete key is pressed to ensure proper behavior.
        /// </summary>
        /// <returns>
        /// Returns true if this method handled the delete key press.
        /// </returns>
        private bool HandleKeyDelete()
        {
            // If the caret is at position 0 and the RichTextEdit control deletes the first
            // character in the _commandLine, it will combine the _promptRun and _commandLine
            // into a single run.
            if ((CommandCaretPosition == 0) && (CommandText.Length > 0))
            {
                _commandLine.Text = _commandLine.Text.Substring(1);
                CaretPosition = _commandLine.ContentStart;
                return true;
            }
            return false;
        }

        /// <summary>
        /// Execute the current PowerShell command text and then display a new prompt.
        /// </summary>
        /// <returns>Always returns true. The RichTextBox is never allowed to process the Enter key.</returns>
        private bool HandleKeyEnter()
        {
            Document.Blocks.Remove(_paragraph);
            this.WriteLine(_prompt + CommandText);

            try
            {
                // Execute the command with no input.
                this.ExecuteCommand(CommandText, null);
            }
            catch (RuntimeException rte)
            {
                this.ReportException(rte);
            }

            NewCommandPrompt();

            return true;
        }

        /// <summary>
        /// Handle the left arrow key press if the caret is at the start of the command text.
        /// </summary>
        /// <returns>Returns true if the caret is at the start of the command text.</returns>
        private bool HandleKeyLeft()
        {
            return CommandCaretAtContentStart;
        }

        /// <summary>
        /// Handle the right arrow key press if the caret is at the end of the command text.
        /// </summary>
        /// <returns>Returns true if the caret is at the end of the command text.</returns>
        private bool HandleKeyRight()
        {
            return CommandCaretAtContentEnd;
        }

        #endregion Key Handler Functions

        /// <summary>
        /// Create a new command prompt in a new document.
        /// </summary>
        /// <param name="NewDocument"></param>
        private void NewCommandPrompt(bool NewDocument = false)
        {
            // Clear all inlines from the paragraph and add the prompt and a blank
            // command line.
            _paragraph = new Paragraph();
            _paragraph.Inlines.Add(_promptRun);
            _commandLine.Text = String.Empty;
            _paragraph.Inlines.Add(_commandLine);
            if (NewDocument)
            {
                Document = new FlowDocument(_paragraph);
                Style style = new Style(typeof(Paragraph));
                style.Setters.Add(new Setter(Block.MarginProperty, new Thickness(0)));
                Document.Resources.Add(typeof(Paragraph), style);
            }
            else
            {
                Document.Blocks.InsertAfter(Document.Blocks.Last(), _paragraph);
            }
            CaretPosition = _commandLine.ContentStart;
        }

        /// <summary>
        /// An exception occurred that we want to display using the display formatter. 
        /// To do this we run a second pipeline passing in the error record. The runtime 
        /// will bind this to the $input variable which is why $input is being piped to 
        /// out-string. We then call WriteErrorLine to make sure the error gets displayed 
        /// in the correct error color.
        /// </summary>
        /// <param name="e">
        /// The exception to display
        /// </param>
        private void ReportException(RuntimeException rte)
        {
            if (rte != null)
            {
                object error;
                IContainsErrorRecord icer = rte as IContainsErrorRecord;
                if (icer != null)
                {
                    error = icer.ErrorRecord;
                }
                else
                {
                    error = (object)new ErrorRecord(rte, "Host.ReportException", ErrorCategory.NotSpecified, null);
                }

                lock (this.instanceLock)
                {
                    this.currentPowerShell = PowerShell.Create();
                }

                this.currentPowerShell.Runspace = this.powerShellRunspace;

                try
                {
                    this.currentPowerShell.AddScript("$input").AddCommand("out-string");

                    // Do not merge errors, this function will swallow errors.
                    Collection<PSObject> result;
                    PSDataCollection<object> inputCollection = new PSDataCollection<object>();
                    inputCollection.Add(error);
                    inputCollection.Complete();
                    result = this.currentPowerShell.Invoke(inputCollection);

                    if (result.Count > 0)
                    {
                        string str = result[0].BaseObject as string;
                        if (!string.IsNullOrEmpty(str))
                        {
                            // Remve \r\n that is added by Out-string.
                            this.powerShellHost.UI.WriteErrorLine(str.Substring(0, str.Length - 2));
                        }
                    }
                }
                finally
                {
                    // Dispose of the pipeline line and set it to null, locked because currentPowerShell
                    // may be accessed by the ctrl-C handler.
                    lock (this.instanceLock)
                    {
                        this.currentPowerShell.Dispose();
                        this.currentPowerShell = null;
                    }
                }
            }
        }

        public void SetBufferSize(int Width, int Height)
        {
            if (buffer == null)
            {
                buffer = new BufferCell[Width, Height];
            }
            else
            {
                BufferCell[,] newBuffer = new BufferCell[Width, Height];
                int minHeight = Math.Min(Width, buffer.GetLength(1));
                int minWidth = Math.Min(Height, buffer.GetLength(0));
                for (int i = 0; i < minWidth; i++)
                {
                    for (int j = 0; j < minHeight; j++)
                    {
                        newBuffer[i, j] = buffer[i, j];
                    }
                }
                buffer = newBuffer;
            }
        }

        public delegate void WriteLineDelegate(ConsoleColor foregroundColor, ConsoleColor backgroundColor, string value);

        public async void WriteLine(string value)
        {
            await this.Dispatcher.BeginInvoke((WriteLineDelegate)this.WriteLineOnUIThread, DispatcherPriority.Normal, this.ForegroundColor, this.BackgroundColor, value);
        }

        public async void WriteLine(ConsoleColor foregroundColor, ConsoleColor backgroundColor, string value)
        {
            await this.Dispatcher.BeginInvoke((WriteLineDelegate)this.WriteLineOnUIThread, DispatcherPriority.Normal, foregroundColor, backgroundColor, value);
        }

        protected void WriteLineOnUIThread(ConsoleColor foregroundColor, ConsoleColor backgroundColor, string value)
        {
            Paragraph paragraph = new Paragraph();
            Run currentRun = new Run(value);
            currentRun.Background = GetBrushFromConsoleColor(backgroundColor);
            currentRun.Foreground = GetBrushFromConsoleColor(foregroundColor);
            paragraph.Inlines.Add(currentRun);
            Document.Blocks.InsertAfter(Document.Blocks.Last(), paragraph);
        }
    }
}
